"""
图形识别系统 - Flask主程序
"""
import os
import cv2
import numpy as np
from PIL import Image
from flask import Flask, render_template, request, jsonify, send_from_directory
from werkzeug.utils import secure_filename
import base64
from io import BytesIO
import json

from core.shape_detector import ShapeDetector
from core.yolo_detector import YoloDetector
from core.password_lock import PasswordLock
from config.settings import DEFAULT_PASSWORD_SEQUENCE, DETECTION_MODE_CONFIG, CAMERA_CONFIG

# 创建Flask应用
template_dir = os.path.abspath('templates')
static_dir = os.path.abspath('static')
app = Flask(__name__, template_folder=template_dir, static_folder=static_dir)

app.config['SECRET_KEY'] = 'shape_recognition_secret_key'
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 最大16MB上传限制

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 初始化核心组件
traditional_detector = ShapeDetector()
yolo_detector = YoloDetector()
password_lock = PasswordLock(DEFAULT_PASSWORD_SEQUENCE)

# 当前检测模式和调试模式
current_detection_mode = DETECTION_MODE_CONFIG["default_mode"]
debug_mode = DETECTION_MODE_CONFIG["debug_mode"]

def pil_to_cv2(pil_image):
    """将PIL图像转换为OpenCV格式"""
    # 转换PIL图像为numpy数组
    numpy_image = np.array(pil_image)
    # 转换RGB到BGR
    cv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
    return cv_image

def cv2_to_pil(cv_image):
    """将OpenCV图像转换为PIL格式"""
    # 转换BGR到RGB
    rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
    # 转换为PIL图像
    pil_image = Image.fromarray(rgb_image)
    return pil_image

def image_to_base64(image):
    """将PIL图像转换为base64编码字符串"""
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str

def detect_shapes_with_mode(image, mode=None):
    """根据指定模式检测图形"""
    global current_detection_mode
    if mode is not None:
        current_detection_mode = mode
    
    if current_detection_mode == "yolo" and yolo_detector is not None:
        return yolo_detector.detect_shapes(image)
    else:
        return traditional_detector.detect_shapes(image)

@app.route('/')
def index():
    """主页路由"""
    return render_template('index.html', 
                         detection_mode=current_detection_mode,
                         debug_mode=debug_mode,
                         camera_enabled=CAMERA_CONFIG["enable_camera"])

@app.route('/detect_shapes', methods=['POST'])
def detect_shapes():
    """检测上传图像中的形状"""
    try:
        # 检查是否有文件上传
        if 'image' not in request.files:
            return jsonify({'error': '没有上传文件'}), 400
        
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': '没有选择文件'}), 400
        
        # 处理上传的图像
        image = Image.open(file)
        cv_image = pil_to_cv2(image)
        
        # 获取检测模式（如果提供了的话）
        mode = request.form.get('mode', current_detection_mode)
        
        # 识别图形
        result_image, shapes = detect_shapes_with_mode(cv_image, mode)
        result_pil = cv2_to_pil(result_image)
        
        # 转换结果图像为base64
        result_base64 = image_to_base64(result_pil)
        
        # 生成结果文本
        if shapes:
            result_text = f"识别到 {len(shapes)} 个图形:\n"
            for i, shape in enumerate(shapes, 1):
                result_text += f"{i}. {shape['name']} (置信度: {shape['confidence']:.2f})\n"
        else:
            result_text = "未识别到任何图形"
        
        # 如果是调试模式，添加额外的调试信息
        if debug_mode:
            result_text += f"\n[调试信息]\n检测模式: {current_detection_mode}\n"
            result_text += f"图像尺寸: {cv_image.shape}\n"
        
        return jsonify({
            'image': result_base64,
            'text': result_text,
            'shapes': shapes,
            'mode': current_detection_mode
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/set_mode', methods=['POST'])
def set_mode():
    """设置检测模式"""
    global current_detection_mode, debug_mode
    try:
        data = request.get_json()
        if 'mode' in data:
            current_detection_mode = data['mode']
        if 'debug' in data:
            debug_mode = data['debug']
        
        return jsonify({
            'success': True,
            'mode': current_detection_mode,
            'debug': debug_mode
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/get_status', methods=['GET'])
def get_status():
    """获取当前系统状态"""
    return jsonify({
        'mode': current_detection_mode,
        'debug': debug_mode,
        'camera_enabled': CAMERA_CONFIG["enable_camera"]
    })

@app.route('/password_status', methods=['GET'])
def password_status():
    """获取密码锁状态"""
    try:
        status = password_lock.get_status()
        return jsonify(status)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/check_password', methods=['POST'])
def check_password():
    """检查密码图形"""
    try:
        # 检查是否有文件上传
        if 'image' not in request.files:
            return jsonify({'error': '没有上传文件'}), 400
        
        file = request.files['image']
        if file.filename == '':
            return jsonify({'error': '没有选择文件'}), 400
        
        # 处理上传的图像
        image = Image.open(file)
        cv_image = pil_to_cv2(image)
        
        # 识别图形
        result_image, shapes = detect_shapes_with_mode(cv_image)
        result_pil = cv2_to_pil(result_image)
        
        # 转换结果图像为base64
        result_base64 = image_to_base64(result_pil)
        
        # 检查密码
        is_correct, is_complete, message = password_lock.check_shapes(shapes)
        
        return jsonify({
            'image': result_base64,
            'message': message,
            'is_correct': is_correct,
            'is_complete': is_complete,
            'shapes': shapes
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/reset_password', methods=['POST'])
def reset_password():
    """重置密码锁"""
    try:
        password_lock.reset()
        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # 从环境变量获取配置，如果没有则使用默认值
    host = os.environ.get('FLASK_HOST', '127.0.0.1')
    port = int(os.environ.get('FLASK_PORT', 5000))
    debug = os.environ.get('FLASK_DEBUG', 'True').lower() == 'true'
    
    print(f"启动图形识别系统...")
    print(f"访问地址: http://{host}:{port}")
    print(f"调试模式: {debug}")
    print(f"检测模式: {current_detection_mode}")
    
    app.run(host=host, port=port, debug=debug)